from llmtcg.chat.models.openai import OpenaiCompatibleChatModel

from typing import Any, Dict, Tuple
from os import environ
from dataclasses import dataclass
from enum import Enum


@dataclass
class BaichuanChatModel(OpenaiCompatibleChatModel):

    """ Models that are deployed at Infini-AI. You need to set the environment
    variable INFINIAI_API_KEY or manually provide the api key when initialize
    the model instance. """

    class Model(Enum):
        Baichuan2_Turbo = "Baichuan2-Turbo"
        Baichuan2_Turbo_192K = "Baichuan2-Turbo-192k"
        Baichuan2_53B = "Baichuan2-53B"

    api_key: str = environ.get("BAICHUAN_API_KEY", "")
    model: Model = Model.Baichuan2_Turbo_192K

    @property
    def api_address(self) -> str:
        return "https://api.baichuan-ai.com/v1"

    @property
    def headers(self) -> Dict[str, str]:
        return {"Authorization": f"Bearer {self.api_key}"}

    @property
    def model_name(self) -> str:
        return self.model.value

    def detect_error(self, response: Any) -> Tuple[bool, str]:
        if 'code' in response:
            if response['code'] == 10003:
                return True, f"authentication error, please check API_KEY (current value is {self.api_key})"

        return False, ""

    def __hash__(self) -> int:
        return id(self)
